# -*- coding: utf-8 -*-

"""
DB Connection

host=None, user=None, password="",
database=None, port=0, unix_socket=None,
charset='', sql_mode=None,
read_default_file=None, conv=None, use_unicode=None,
client_flag=0, cursorclass=Cursor, init_command=None,
connect_timeout=10, ssl=None, read_default_group=None,
compress=None, named_pipe=None,
autocommit=False, db=None, passwd=None, local_infile=False,
max_allowed_packet=16*1024*1024, defer_connect=False,
auth_plugin_map=None, read_timeout=None, write_timeout=None,
bind_address=None, binary_prefix=False, program_name=None,
server_public_key=None
"""

from __future__ import absolute_import, print_function

import copy
import inspect
import logging
import threading
import time
import uuid

import pymysql
from six.moves.queue import Queue

from lian.orm import statement
from lian.utils.naming import camel2underline

DEFAULT_DB = 'default'
DEFAULT_CONNECTIONS = 5
DEFAULT_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor,
    'MaxConnections': DEFAULT_CONNECTIONS,
}
LOGGER_NAME = 'lian.orm.db'
LOG = logging.getLogger(LOGGER_NAME)
VALID_DB_PARAMS = inspect.getargs(pymysql.connections.Connection.__init__.__code__).args[1:]


class PooledConnection(object):
    """连接"""

    def __init__(self, pool, db_label, connection):
        self._connection = connection
        self._db_label = db_label
        self._pool = pool
        self.logger = pool.get_logger(db_label)

    def close(self):
        if self._connection is not None:
            if self._pool.is_connection_using(self._connection, self._db_label):
                self._pool.release(self._connection, self._db_label)
            self._connection = None

    def __getattr__(self, name):
        # All other members are the same.
        return getattr(self._connection, name)

    def __del__(self):
        self.close()


def init(config):
    ConnectionPool.init(config)


class ConnectionPool(object):
    """连接池"""
    _instance_lock = threading.Lock()

    _config = {}
    _config_grouped = {}
    logger = None
    default_db = DEFAULT_DB

    @classmethod
    def init(cls, config):
        if not cls.logger:
            cls.logger = LOG
        cls.logger.debug('ConnectionPool.init')
        if cls.is_inited():
            raise Exception('ConnectionPool inited already')

        assert isinstance(config, dict)
        for db_label, db_config in config.items():
            assert isinstance(db_label, str)
            assert isinstance(db_config, dict)
            if db_config.get('disabled', False):
                cls.logger.info('%s: disabled', db_label)
                continue
            db_config.update({k: copy.deepcopy(v) for k, v in DEFAULT_CONFIG.items() if k not in db_config})
            db_config['db_label'] = db_label
            if 'database' not in db_config:
                db_config['database'] = db_config.get('db', db_label)
            cls._config[db_label] = db_config
            cls._config_grouped.setdefault(db_config['database'], []).append(db_label)

        if cls.default_db not in cls._config:
            db_labels = list(cls._config.keys())
            cls.default_db = db_labels[0]
            cls.logger.debug('All databases: %r', db_labels)
            cls.logger.info('The database config has not DEFAULT_DB, choose %s as default database', cls.default_db)

    @classmethod
    def is_inited(cls):
        return cls._config != {}

    @classmethod
    def ensure_inited(cls):
        if not cls.is_inited():
            raise Exception('ConnectionPool need initialize')

    @classmethod
    def real_db(cls, db_label=None):
        if not db_label or db_label == DEFAULT_DB:
            db_label = cls.default_db
        return db_label

    @classmethod
    def db_labels(cls, db_name):
        return cls._config_grouped.get(db_name, [])

    @classmethod
    def set_logger(cls, logger):
        cls.logger = logger

    @classmethod
    def get_config(cls, db_label=None):
        return cls._config[cls.real_db(db_label)]

    def get_logger(self, db_label=None):
        return self._loggers.get(self.real_db(db_label), self.logger)

    def get_queue(self, db_label=None):
        return self._queues[self.real_db(db_label)]

    def get_queue_using(self, db_label=None):
        return self._queues_using[self.real_db(db_label)]

    def __init__(self):
        self.ensure_inited()
        self.logger.info('init')

        self._queues = {}
        self._queues_using = {}
        self._loggers = {}

        for db_label in self._config:
            db_config = self.get_config(db_label)

            logger = db_config.pop('Logger', None)
            logger = logger or logging.getLogger('%s:%s' % (LOGGER_NAME, db_label))
            assert isinstance(logger, logging.Logger), 'param logger of db config error: %r' % logger
            self._loggers[db_label] = logger

            _max_connections = db_config.pop('MaxConnections', DEFAULT_CONNECTIONS)
            assert isinstance(_max_connections, int) and _max_connections > 0

            self._queues[db_label] = Queue(_max_connections)  # create the queue
            self._queues_using[db_label] = []

            self.build_connection(db_label)  # build a connection

    def get_queue_status(self, db_label):
        db_label = self.real_db(db_label)
        q = self.get_queue(db_label)
        q_using = self.get_queue_using(db_label)
        return 'queue of %s size: %s + %s / %s' % (db_label, q.qsize(), len(q_using), q.maxsize)

    def is_connection_using(self, conn, db_label):
        logger = self.get_logger(db_label)
        if db_label is None:
            return any((self.is_connection_using(conn, db_label) for db_label in self._config))
        db_label = self.real_db(db_label)
        id_conn = id(conn)
        q_using = self.get_queue_using(db_label)
        logger.debug('is_connection_using(%s): %r in %r', db_label, id_conn, q_using)
        return id_conn in q_using

    def build_connection(self, db_label):
        logger = self.get_logger(db_label)
        db_label = self.real_db(db_label)
        q = self.get_queue(db_label)
        q_using = self.get_queue_using(db_label)
        if q.qsize() + len(q_using) >= q.maxsize:
            return
        logger.debug('build connection at %s', db_label)
        conn = self.connect(db_label)
        q.put(conn)
        logger.debug('after build connection, %s', self.get_queue_status(db_label))

    @classmethod
    def instance(cls):
        cls.ensure_inited()
        if not hasattr(cls, '_instance'):
            with cls._instance_lock:
                if not hasattr(cls, '_instance'):
                    cls._instance = cls()
        return cls._instance

    def connect(self, db_label):
        logger = self.get_logger(db_label)
        self.ensure_inited()

        db_label = self.real_db(db_label)
        config = self.get_config(db_label)
        _config = copy.deepcopy(config)
        if 'passwd' in _config:
            _config['passwd'] = '*' * 6
        if 'password' in _config:
            _config['password'] = '*' * 6
        logger.debug('connecting database: %r', _config)

        while True:
            try:
                params = {k: v for k, v in config.items() if k in VALID_DB_PARAMS}
                conn = pymysql.connect(**params)
                logger.debug('ping...')
                conn.ping()
                break
            except Exception as e:
                logger.exception('connect failed: %s, retrying...', e)
                time.sleep(1)

        logger.info('database connected')
        return conn

    def acquire(self, db_label):
        logger = self.get_logger(db_label)
        self.ensure_inited()

        db_label = self.real_db(db_label)
        q = self.get_queue(db_label)
        q_using = self.get_queue_using(db_label)
        logger.debug('acquire connection at %s', db_label)
        if q.empty():
            self.build_connection(db_label)

        while True:
            conn = q.get()
            conn_error = False
            try:
                logger.debug('ping...')
                conn.ping()
                break
            except Exception as e:
                logger.exception(e)
                conn_error = True

            if conn_error:
                logger.warning('Connection error, try to close it...')
                try:
                    conn.close()
                except Exception as conn_close_error:
                    logger.exception(conn_close_error)
                self.build_connection(db_label)

        q_using.append(id(conn))
        logger.debug('after acquire connection, %s', self.get_queue_status(db_label))
        return conn

    def release(self, conn, db_label):
        logger = self.get_logger(db_label)
        # import traceback
        # logger.debug('TraceBack:\n' + (''.join(traceback.format_stack())))

        self.ensure_inited()

        db_label = self.real_db(db_label)
        id_conn = id(conn)

        q = self.get_queue(db_label)
        q_using = self.get_queue_using(db_label)

        if not self.is_connection_using(conn, db_label):
            logger.warning('The connection #%d is not using, ignore release... (using connections: %r)',
                           id_conn, q_using)
            return

        logger.debug('release connection %d at %s', id_conn, db_label)
        q.put(conn)
        q_using.remove(id_conn)
        logger.debug('after release connection, %s', self.get_queue_status(db_label))


class ConnectionContext:
    def __init__(self, db_label):
        self.db_label = db_label
        self.connection = None  # pymysql.Connection

    def __enter__(self):
        pool = ConnectionPool.instance()
        self.connection = pool.acquire(self.db_label)
        return PooledConnection(pool, self.db_label, self.connection)

    def __exit__(self, *args):
        """args: type, value, trace"""
        pool = ConnectionPool.instance()
        if pool.is_connection_using(self.connection, self.db_label):
            pool.release(self.connection, self.db_label)


def _execute(sql, need_return=False, auto_commit=False, db_label=None):
    """Execute SQL

    :param sql:
    :param need_return:
    :param auto_commit:
    :param db_label:
    :return:
    """
    with ConnectionContext(db_label) as conn:
        query_uuid = uuid.uuid1()
        conn.logger.debug('[%s] db %s: sql execute start...', query_uuid, db_label)
        conn.logger.debug('[%s] db %s: %s', query_uuid, db_label, sql)
        cur = None
        started_at = time.time()
        try:
            cur = conn.cursor()
            rowcount = cur.execute(sql)
            if auto_commit:
                conn.commit()

            if need_return:
                result = {'rows': cur.fetchall()}
            else:
                result = {'rows': None}

            result.update({
                'conn': conn,
                'rowcount': rowcount,
                'description': cur.description,
                'lastrowid': cur.lastrowid,
            })
            # conn.logger.debug('[%s] %r', query_uuid, result)
            return result
        except Exception as e:
            conn.logger.exception('[%s] db_label %s: %s', query_uuid, db_label, e)
        finally:
            time_cost = time.time() - started_at
            if time_cost > 100:
                conn.logger.warning('[%s] slow sql, cost: %f', query_uuid, time_cost)
            else:
                conn.logger.debug('[%s] cost: %f', query_uuid, time_cost)
            if cur:
                conn.logger.debug('[%s] close cursor...', query_uuid)
                cur.close()
            conn.logger.debug('[%s] db_label %s: sql execute over...', query_uuid, db_label)


def execute(sql, auto_commit=False, db_label=None):
    return _execute(sql, need_return=False, auto_commit=auto_commit, db_label=db_label)


def query(sql, auto_commit=True, db_label=None):
    return _execute(sql, need_return=True, auto_commit=auto_commit, db_label=db_label)


class ObjectNotFound(Exception):
    pass


class BASE(object):
    __database__ = DEFAULT_DB  # 真实的数据库名称或 default
    __table__ = ''  # 表名
    __pk__ = 'id'  # 主键
    __fields__ = tuple()

    def __init__(self):
        pool = ConnectionPool.instance()
        if self.__database__ == DEFAULT_DB:
            self.database_name = pool.get_config(pool.default_db)['database']
        else:
            self.database_name = self.__database__
        self.logger = pool.logger
        self.sql = statement.SQL(self.table_name, database=self.database_name, logger=self.logger)

    @property
    def table_name(self):
        if self.__class__ is BASE:
            raise NotImplementedError
        if self.__table__:
            return self.__table__
        return camel2underline(self.__class__.__name__)

    @property
    def full_table_name(self):
        return '%s.%s' % (self.database_name, self.table_name)

    def get_logger(self, db_label):
        pool = ConnectionPool.instance()
        return pool.get_logger(db_label)

    def get(self, pk, key=None):
        if not key:
            key = self.__pk__
        rows = self.select(conditions={key: pk})
        if not rows:
            raise ObjectNotFound('%s #%s' % (self.full_table_name, pk))
        return rows[0]

    def select(self, fields=None, conditions=None, limit=None, offset=None, order_by=None, group_by=None, raw_sql=None, flag=None):
        if not flag:
            flag = {'table': self.table_name, 'action': 'select'}
        sql = raw_sql or self.sql.select(fields, conditions, limit, offset, order_by, group_by)
        result = query(sql, db_label=self.strategy(flag))
        return result['rows'] if result else []

    def find(self, fields=None, conditions=None, offset=None, order_by=None, raw_sql=None):
        count = self.count(conditions=conditions)
        if count == 0:
            return None
        rows = self.select(fields, conditions, limit=1, offset=offset, order_by=order_by, raw_sql=raw_sql)
        return rows[0]

    def insert(self, values, fields=None, mode='insert',
               update=None,  # mode = update
               conditions=None,  # mode = insert-not-exists
               refetch=False,
               flag=None):
        if not fields:
            fields = self.__fields__
        if not flag:
            flag = {'table': self.table_name, 'action': 'insert'}
        db_label = self.strategy(flag)
        logger = self.get_logger(db_label)
        sql = self.sql.insert(values, fields=fields, mode=mode, update=update, conditions=conditions)
        result = execute(sql, auto_commit=True, db_label=db_label)
        if not result:
            logger.warning('insert return %s: %s', result, sql)
            return None
        return self.get(result['lastrowid']) if refetch else result['lastrowid']

    def insert_many(self, fields, values_list, update_fields=None, flag=None):
        if not fields:
            fields = self.__fields__
        if not flag:
            flag = {'table': self.table_name, 'action': 'insert'}
        db_label = self.strategy(flag)
        logger = self.get_logger(db_label)
        sql = self.sql.insert_many(fields, values_list, update_fields)
        result = execute(sql, auto_commit=True, db_label=db_label)
        if not result:
            logger.warning('insert return %s: %s', result, sql)
            return None

        # {
        #     'rows': None,
        #     'conn': <lian.orm.db.PooledConnection object at 0x7fae315a17f0>,
        #     'rowcount': 500,
        #     'description': None,
        #     'lastrowid': 316613
        # }
        excepted_rows = len(values_list)
        if result['rowcount'] != excepted_rows:
            logger.warning('insert %s rows (excepted: %s)', result['rowcount'], excepted_rows)
        return result

    def update(self, values, conditions=None, flag=None):
        if not flag:
            flag = {'table': self.table_name, 'action': 'update'}
        db_label = self.strategy(flag)
        sql = self.sql.update(values, conditions=conditions)
        result = execute(sql, auto_commit=True, db_label=db_label)
        return result['rowcount']  # 影响行数

    def count(self, conditions=None, flag=None):
        if not flag:
            flag = {'table': self.table_name, 'action': 'count'}
        db_label = self.strategy(flag)
        sql = self.sql.count(conditions)
        result = query(sql, db_label=db_label)
        return result['rows'][0]['COUNT(1)'] if result else 0

    def delete(self, conditions=None, flag=None):
        if not flag:
            flag = {'table': self.table_name, 'action': 'delete'}
        db_label = self.strategy(flag)
        sql = self.sql.delete(conditions)
        result = execute(sql, auto_commit=True, db_label=db_label)
        return result['rowcount']  # 影响行数

    def strategy(self, flag=None):
        pool = ConnectionPool.instance()
        chosen = None
        db_labels = pool.db_labels(self.database_name)
        assert db_labels, 'no config is availabled for %s' % self.database_name
        if not flag:
            return db_labels[0]
        for db_label in db_labels:
            config = pool.get_config(db_label)
            tables = config.get('tables', None)
            if tables and flag['table'] not in tables:
                self.logger.debug('%s: denied by table restriction: %r, %r',
                                  db_label, flag['table'], tables)
                continue
            actions = config.get('actions', None)
            if actions and flag['action'] not in actions:
                self.logger.debug('%s: denied by action restriction: %r, %r',
                                  db_label, flag['action'], actions)
                continue
            chosen = db_label
            break
        assert chosen, 'no config matched for %s' % self.database_name
        self.logger.debug('choose %s', chosen)
        return chosen
